from typing import List

from tree_rag.session import Message

from . import DEFAULT_MODEL_NAME, Prompt, client

# 使用大模型预训练知识回答
SYSTEM_PROMPT = """{role_prompt}"""


PROMPT = """这是你与用户的对话记录：
{history}

你的任务是使用你的预训练知识来回答用户的问题。由于没有检索到相关的知识库内容，你需要基于你的预训练知识来回答。
请注意：
1. 回答要准确、客观，不要编造信息
2. 如果你不确定答案，可以坦诚表示不确定
3. 回答要简洁明了，直接解答用户的问题
4. 回答的开头必须添加提示语："您好！您咨询的问题目前不在我的知识库范围内，以下答案仅供参考，具体请咨询当地经办机构！"

user: {user_say}
"""


def build_prompt(
    user_say: str,
    histories: List[Message],
    role_prompt: str,
) -> "Prompt":
    system_prompt = SYSTEM_PROMPT.format(role_prompt=role_prompt)
    history = "\n".join([item.to_str() for item in histories])

    user_prompt = PROMPT.format(history=history, user_say=user_say)
    return Prompt(system_prompt=system_prompt, user_prompt=user_prompt)


async def run_agent(
    user_say: str,
    histories: List[Message],
    role_prompt: str,
) -> str:
    """Run the agent.

    Args:
        user_say: The user's input.
        histories: The conversation history.
        role_prompt: The role prompt.
    """
    prompt: Prompt = build_prompt(user_say, histories, role_prompt)
    chat_completion = await client.chat.completions.create(
        messages=[
            {"role": "system", "content": prompt.system_prompt},
            {"role": "user", "content": prompt.user_prompt},
        ],
        model=DEFAULT_MODEL_NAME,
    )
    response = chat_completion.choices[0].message.content
    return response 